import pickle
import gradio as gr
import numpy as np
from PIL import Image
from sklearn.preprocessing import MinMaxScaler
from pinecone import Pinecone, ServerlessSpec

# 创建Pinecone实例
pc = Pinecone(
    api_key='816f2326-8a49-460e-8818-d810d393d713'  # 替换为密钥
)

# 选择环境
environment = 'us-east-1'  # 替换为环境

# 连接到索引
index_name = 'mnist-index'  # 替换为索引名称
if index_name not in pc.list_indexes().names():
    # 如果索引不存在，则创建索引
    pc.create_index(
        name=index_name,
        dimension=64,  # 假设每个向量是8x8图像，所以维度是64
        metric='euclidean',
        pod_type='p1',  # 或者你选择的Pod类型
        metadata_config={'indexed': ['label']}  # 如果你有元数据，可以在这里配置
    )
index = pc.Index(index_name)

# 定义预测函数，这个函数将用于Gradio接口进行预测
def predict_digit(image):
    # 预处理图像
    image = image.convert('L')
    image = image.resize((8, 8))
    image_array = np.array(image).flatten()
    image_array = (image_array / 255.0).reshape(1, -1)

    # 打印预处理后的图像数组
    print("Preprocessed image array:", image_array)

    # 使用Pinecone进行向量搜索
    query_results = index.query(queries=[image_array.tolist()], top_k=1)
    print("Query results:", query_results)  # 打印查询结果

    # 确保查询结果不为空
    if query_results['matches']:
        # 获取最相似的向量ID
        prediction = query_results['matches'][0]['id']
        return int(prediction)
    else:
        # 如果查询结果为空，返回错误信息
        return "No matches found"
# 创建Gradio接口
with gr.Blocks() as app:
    with gr.Row():
        sketchpad = gr.Sketchpad(label="绘制数字", type="pil")
        output_label = gr.Label(label="预测结果")
    
    with gr.Row():
        submit_button = gr.Button("预测数字")
        # 当用户点击预测按钮时，调用 predict_digit 函数进行预测
    submit_button.click(
        fn=predict_digit,
        inputs=sketchpad,
        outputs=output_label
    )

# 启动Gradio接口，不创建公共链接
app.launch( )  # 设置share=False以启动本地服务器

